/*
 * hio_net.c - TCP/IP networking access to PS2 HDD
 * $Id: hio_udpnet.c,v 1.4 2006/09/01 17:27:02 bobi Exp $
 *
 * Copyright 2004 Bobi B., w1zard0f07@yahoo.com
 *
 * This file is part of hdl_dump.
 *
 * hdl_dump is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * hdl_dump is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with hdl_dump; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#if defined (_BUILD_WIN32)
#  if defined (_MSC_VER) && defined (_WIN32)
#    include <winsock2.h> /* Microsoft Visual C/C++ compiler */
#  else
#    include <winsock.h> /* GNU C/C++ compiler */
#  endif
#  include <windows.h>
#elif defined (_BUILD_UNIX)
#  include <errno.h>
#  include <unistd.h>
#  include <sys/types.h>
#  include <sys/socket.h>
#  include <netinet/in.h>
#  include <netinet/tcp.h>
#endif
#include <assert.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include "osal.h"
#include "hio_udpnet.h"
#include "byteseq.h"
#include "net_common.h"
#include "retcodes.h"
#include "progress.h"
#include "net_io.h"


typedef struct hio_net_type
{
  hio_t hio;
  SOCKET sock;
  SOCKET udp;
  unsigned long error_code;
  size_t target_kbps;
  int auto_throttle;
} hio_net_t;

#define SETBIT(mask, bit) (mask)[(bit) / 32] |= 1 << ((bit) % 32)
#define GETBIT(mask, bit) ((mask)[(bit) / 32] & (1 << ((bit) % 32)))

#if defined (_BUILD_WIN32)
static int net_init = 0;
#endif

#if defined (_BUILD_WIN32)
/* it is not the same, since usleep works in microseconds,
 * whereas Sleep is in miliseconds */
#  define delay(ms) Sleep (ms)
#else
#  define delay(ms) usleep((ms) * 1000)
#endif


/**************************************************************/
/* execute a command and wait for a reply */
static int
query (SOCKET s,
       u_int32_t command,
       u_int32_t sector,
       u_int32_t num_sectors,
       u_int32_t *response,
       char output [HDD_SECTOR_SIZE * NET_NUM_SECTORS]) /* or NULL */
{
  unsigned char cmd [NET_IO_CMD_LEN];
  u_int32_t cmd_length = NET_IO_CMD_LEN;
  int result;
  set_u32 (cmd +  0, command);
  set_u32 (cmd +  4, sector);
  set_u32 (cmd +  8, num_sectors);
  set_u32 (cmd + 12, 0x003a2d29);

  result = send_exact (s, (const char*) cmd, cmd_length, 0);
  if (result == cmd_length)
    { /* command successfully sent */
      result = recv_exact (s, (char*) cmd, NET_IO_CMD_LEN, 0);
      if (result == NET_IO_CMD_LEN)
	{ /* response successfully received */
	  if (get_u32 (cmd + 0) != command ||
	      get_u32 (cmd + 4) != sector ||
	      get_u32 (cmd + 8) != num_sectors)
	    return (RET_PROTO_ERR);

	  *response = get_u32 (cmd + 12);
	  if (output != NULL &&
	      *response != (u_int32_t) -1)
	    { /* receive additional information */
	      size_t bytes_expected = HDD_SECTOR_SIZE * num_sectors;
	      result = recv_exact (s, output, bytes_expected, 0);
	      result = result == bytes_expected ? RET_OK : RET_ERR;
	    }
	  else
	    result = RET_OK;
	}
      else
	result = RET_ERR;
    }
  else
    result = RET_ERR;
  return (result);
}


/**************************************************************/
#define FEEDBACK 0
static int
send_for_write (hio_net_t *net,
		u_int32_t command,
		u_int32_t sector,
		u_int32_t num_sectors,
		u_int32_t *response,
		const char data [HDD_SECTOR_SIZE * NET_NUM_SECTORS])
{
  static size_t retries = 0, sequence = 0;
  time_t delay_time = 1; /* <= tune here */

  u_int32_t bitmask[(NET_NUM_SECTORS + 31) / 32] = { 0 };
  typedef struct udp_packet_t
  { /* 64-bit fix: command and start were unsigned long */
    unsigned char sector[HDD_SECTOR_SIZE * 2];
    u_int32_t command, start;
  } udp_packet_t;
  udp_packet_t packet;
  int result;

#if (FEEDBACK == 3)
  printf ("\t\t\t\t\t\t\t%uKBps,r%u,q%u\r",
	  (unsigned int) net->target_kbps,
	  (unsigned int) retries,
	  (unsigned int) sequence);
#endif

  /* ask to start writing operation */
  result = query (net->sock, CMD_HIO_WRITE,
		  sector, num_sectors, response, NULL);
  if (result == RET_OK &&
      *response == 0)
    {
      do
	{ /* flood with (unconfirmed) data */
	  size_t sps = net->target_kbps; /* *1032 */
	  size_t i, count = 0;

	  /* new throttling code */
	  highres_time_t start;
	  int sent = 0;

	  for (i = 0; i < num_sectors; i += 2)
	    if (!GETBIT (bitmask, i))
	      ++count;

	  highres_time (&start);
	  while (count)
	    {
	      double seconds; /* curr time relative to start */
	      double next; /* the time when the next packet should be sent */
	      do
		{ /* what's the time, how many seconds have gone and
		   * what is the ratio of job that _has_ to be done
		   * against the job that _is_ done */
		  highres_time_t now;

		  highres_time (&now);
		  seconds = (double) ((highres_time_val (&now) -
				       highres_time_val (&start)) /
				      (double) HIGHRES_TO_SEC);
		  next = (double) sent / (double) sps;
#if (FEEDBACK == 2)
		  printf ("%.5f, %.5f, %4d\n",
			  next, seconds, sent);
		  fflush (stdout);
#endif

		  if (sent == 0 || next <= seconds)
		    { /* lagging behind... */
		      for (i = 0; i < num_sectors; i += 2)
			if (!GETBIT (bitmask, i))
			  {
			    memcpy (packet.sector, data + i * HDD_SECTOR_SIZE,
				    HDD_SECTOR_SIZE * 2);
			    set_u32 (&packet.command, command);
			    set_u32 (&packet.start, sector + i);
			    (void) send (net->udp, (void*) &packet,
					 sizeof (packet), 0);
#if (FEEDBACK == 1)
			    printf (".");
#endif

			    /* send a single packet only */
			    SETBIT (bitmask, i);
			    break;
			  }

		      ++sent;
		      --count;
		    }
		}
	      while (count > 0 && next < seconds); /* quick loop */

	      if (count)
		{
		  delay (delay_time);
#if (FEEDBACK == 1)
		  printf ("z");
#endif
		}
	    } /* while more data waiting to be sent... */

	  /* ask if all data is received */
	  result = query (net->sock, CMD_HIO_WRITE_STAT, sector, num_sectors,
			  response, NULL);
	  if (result == RET_OK)
	    {
	      if (*response == num_sectors)
		{ /* might increase transfer speed */
		  ++sequence;
		  if (net->auto_throttle && (retries < 3 || sequence > 30))
		    {
		      net->target_kbps += 10;
		      sequence = 0;
		    }
		  result = RET_OK;
		  break; /* successfully completed */
		}
	      else if (*response == (u_int32_t) -1)
		result = RET_SVR_ERR; /* write operation has failed */
	      else if (*response == 0)
		{ /* more data needed; bitmask has been sent */
		  /* 64-bit fix: replaced unsigned long with u_int32_t */
		  u_int32_t tmp[(NET_NUM_SECTORS + 31) / 32];
		  result = recv_exact (net->sock, (void*) tmp,
				       sizeof (tmp), 0);
		  if (result == sizeof (tmp))
		    { /* refresh bitmask */
		      for (i = 0; i < (NET_NUM_SECTORS + 31) / 32; ++i)
			bitmask[i] = get_u32 (tmp + i);

		      /* retries necessary => decrease transfer speed */
		      if (net->auto_throttle && net->target_kbps > 10)
			net->target_kbps -= 10;
		      ++retries;
		      sequence = 0;
		      result = RET_OK;
		    }
		  else
		    result = RET_ERR;
		}
	      else
		result = RET_SVR_ERR; /* protocol error? */
	    }
	}
      while (result == RET_OK);
    }
  return (result);
}


/**************************************************************/
/* execute a command with no reply */
static int
execute (SOCKET s,
	 u_int32_t command,
	 u_int32_t sector,
	 u_int32_t num_sectors)
{
  unsigned char cmd [NET_IO_CMD_LEN + HDD_SECTOR_SIZE * NET_NUM_SECTORS];
  u_int32_t cmd_length = NET_IO_CMD_LEN;
  int result;
  set_u32 (cmd +  0, command);
  set_u32 (cmd +  4, sector);
  set_u32 (cmd +  8, num_sectors);
  set_u32 (cmd + 12, 0x7d3a2d29); /* }:-) */

  result = send_exact (s, (const char*) cmd, cmd_length, 0);
  return (result == cmd_length ? RET_OK : RET_ERR);
}


/**************************************************************/
static int
net_stat (hio_t *hio,
	  /*@out@*/ u_int32_t *size_in_kb)
{
  hio_net_t *net = (hio_net_t*) hio;
  u_int32_t size_in_kb2;
  int result = query (net->sock, CMD_HIO_STAT, 0, 0, &size_in_kb2, NULL);
  if (result == OSAL_OK)
    *size_in_kb = size_in_kb2;
  else
    net->error_code = osal_get_last_error_code ();
  return (result);
}


/**************************************************************/
static int
net_read (hio_t *hio,
	  u_int32_t start_sector,
	  u_int32_t num_sectors,
	  /*@out@*/ void *output,
	  /*@out@*/ u_int32_t *bytes)
{
  hio_net_t *net = (hio_net_t*) hio;
  u_int32_t response;
  int result;
  char *outp = (char*) output;

  *bytes = 0;
  do
    { /* TODO: test the maximum read size? */
#if 0
      u_int32_t at_once_s = (num_sectors > NET_NUM_SECTORS ?
			     NET_NUM_SECTORS : num_sectors);
#else
      u_int32_t at_once_s = (num_sectors > 32 ? 32 : num_sectors);
#endif
      result = query (net->sock, CMD_HIO_READ, start_sector, at_once_s,
		      &response, outp);
      if (result == OSAL_OK)
	{
	  if (response == at_once_s)
	    {
	      start_sector += at_once_s;
	      num_sectors -= at_once_s;
	      *bytes += HDD_SECTOR_SIZE * at_once_s;
	      outp += HDD_SECTOR_SIZE * at_once_s;
	    }
	  else
	    /* server reported an error; give up */
	    result = RET_SVR_ERR;
	}
      else
	net->error_code = osal_get_last_error_code ();
    }
  while (result == OSAL_OK && num_sectors > 0);
  return (result);
}


/**************************************************************/
static int
net_write (hio_t *hio,
	   u_int32_t start_sector,
	   u_int32_t num_sectors,
	   const void *input,
	   /*@out@*/ u_int32_t *bytes)
{
  hio_net_t *net = (hio_net_t*) hio;
  u_int32_t response;
  int result = RET_OK;
  char *inp = (char*) input;

  if (result == RET_OK)
    {
      *bytes = 0;
      do
	{
	  u_int32_t at_once_s = (num_sectors > NET_NUM_SECTORS ?
				 NET_NUM_SECTORS : num_sectors);
	  const void *data_to_send;
	  u_int32_t sectors_to_send;

	  data_to_send = inp;
	  sectors_to_send = at_once_s;

	  result = send_for_write (net, CMD_HIO_WRITE,
				   start_sector, sectors_to_send,
				   &response, data_to_send);
	  if (result == OSAL_OK)
	    {
	      if (response == at_once_s)
		{
		  start_sector += at_once_s;
		  num_sectors -= at_once_s;
		  *bytes += HDD_SECTOR_SIZE * at_once_s;
		  inp += HDD_SECTOR_SIZE * at_once_s;
		}
	      else
		/* server reported an error; give up */
		result = RET_SVR_ERR;
	    }
	  else
	    net->error_code = osal_get_last_error_code ();
	}
      while (result == OSAL_OK && num_sectors > 0);
    }
  return (result);
}


/**************************************************************/
static int
net_poweroff (hio_t *hio)
{
  hio_net_t *net = (hio_net_t*) hio;
  int result = execute (net->sock, CMD_HIO_POWEROFF, 0, 0);
  return (result);
}


/**************************************************************/
static int
net_flush (hio_t *hio)
{
  hio_net_t *net = (hio_net_t*) hio;
  u_int32_t response;
  int result = query (net->sock, CMD_HIO_FLUSH, 0, 0, &response, NULL);
  return (result);
}


/**************************************************************/
static int
net_close (/*@special@*/ /*@only@*/ hio_t *hio) /*@releases hio@*/
{
  hio_net_t *net = (hio_net_t*) hio;
#if defined (_BUILD_WIN32)
  shutdown (net->sock, SD_RECEIVE | SD_SEND);
  closesocket (net->sock);
  closesocket (net->udp);
#elif defined (_BUILD_UNIX)
  close (net->sock);
  close (net->udp);
#endif
  osal_free (hio);

#if defined (_BUILD_WIN32)
  timeEndPeriod (1);
#endif

  return (RET_OK);
}


/**************************************************************/
static char*
net_last_error (hio_t *hio)
{
  hio_net_t *net = (hio_net_t*) hio;
  return (osal_get_error_msg (net->error_code));
}


/**************************************************************/
static void
net_dispose_error (hio_t *hio,
		   /*@only@*/ char* error)
{
  osal_dispose_error_msg (error);
}


/**************************************************************/
static hio_t*
net_alloc (const dict_t *config,
	   SOCKET tcp,
	   SOCKET udp)
{
  hio_net_t *net = (hio_net_t*) osal_alloc (sizeof (hio_net_t));
  if (net != NULL)
    {
      memset (net, 0, sizeof (hio_net_t));
      net->hio.stat = &net_stat;
      net->hio.read = &net_read;
      net->hio.write = &net_write;
      net->hio.flush = &net_flush;
      net->hio.close = &net_close;
      net->hio.poweroff = &net_poweroff;
      net->hio.last_error = &net_last_error;
      net->hio.dispose_error = &net_dispose_error;
      net->sock = tcp;
      net->udp = udp;

      /* initial one here; the rest is done by the auto-tuning code */
      net->target_kbps = dict_get_numeric (config, CONFIG_TARGET_KBPS, 2400);
      net->auto_throttle = dict_get_numeric (config, CONFIG_AUTO_THROTTLE, 0);

#if defined (_BUILD_WIN32)
      timeBeginPeriod (1);
#endif
    }
  return ((hio_t*) net);
}


/**************************************************************/
int
hio_udpnet_probe (const dict_t *config,
		  const char *path,
		  hio_t **hio)
{
  int result = RET_NOT_COMPAT;
  char *endp;
  int a, b, c, d;

#if defined (_BUILD_WIN32)
  /* only Windows requires sockets initialization */
  if (!net_init)
    {
      WORD version = MAKEWORD (2, 2);
      WSADATA wsa_data;
      int result = WSAStartup (version, &wsa_data);
      if (result == 0)
	net_init = 1; /* success */
      else
	return (RET_ERR);
    }
#endif

  a = strtol (path, &endp, 10);
  if (a > 0 && a <= 255 && *endp == '.')
    { /* there is a chance */
      b = strtol (endp + 1, &endp, 10);
      if (b >= 0 && b <= 255 && *endp == '.')
	{
	  c = strtol (endp + 1, &endp, 10);
	  if (c >= 0 && c <= 255 && *endp == '.')
	    {
	      d = strtol (endp + 1, &endp, 10);
	      if (d >= 0 && d <= 255 && *endp == '\0')
		{
		  SOCKET s = socket (PF_INET, SOCK_STREAM, 0);
		  if (s != INVALID_SOCKET)
		    { /* generally there ain't clean-up below, but that is
		       * not really fatal for such class application */
		      struct sockaddr_in sa;
		      memset (&sa, 0, sizeof (sa));
		      sa.sin_family = AF_INET;
		      sa.sin_addr.s_addr = htonl ((a << 24) | (b << 16) | (c << 8) | (d));
		      sa.sin_port = htons (NET_HIO_SERVER_PORT);
		      result = connect (s, (const struct sockaddr*) &sa,
					sizeof (sa)) == 0 ? RET_OK : RET_ERR;
		      if (result == 0)
			{ /* socket connected */
			  SOCKET udp = socket (PF_INET, SOCK_DGRAM, 0);
			  result = connect (udp, (const struct sockaddr*) &sa,
					    sizeof (sa));
			  if (result == 0)
			    {
			      *hio = net_alloc (config, s, udp);
			      if (*hio != NULL)
				; /* success */
			      else
				result = RET_NO_MEM;
			    }
			  else
			    result = RET_ERR;
			}
		      else
			result = RET_ERR;

		      if (result != RET_OK)
			{ /* close socket on error */
#if defined (_BUILD_WIN32)
			  DWORD err = GetLastError ();
			  shutdown (s, SD_RECEIVE | SD_SEND);
			  closesocket (s);
			  SetLastError (err);
#elif defined (_BUILD_UNIX)
			  int err = errno;
			  close (s);
			  errno = err;
#endif
			}
		    }
		}
	    }
	}
    }
  return (result);
}
